package net.csiit.platform.codeguard.agent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtConstructor;
import javassist.CtMethod;
import net.csiit.platform.codeguard.util.StringUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.nio.charset.StandardCharsets;
import java.security.ProtectionDomain;

/**
 * @author adam
 * @date 2020/7/2
 */
public class SpringBootJarLauncherTransformer implements ClassFileTransformer {
    public static final String MAIN_METHOD_RUNNER_SOURCE_FILE_NAME = "/beforeSpringBootStart.txt";
    public static final String SIMPLE_METADATA_READER_SOURCE_FILE_NAME = "/beforeSimpleMetadataReader.txt";

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        ClassPool classPool = ClassPool.getDefault();

        if (StringUtils.equals(className, "org/springframework/boot/loader/MainMethodRunner")) {
            return transformMainMethodRunner(classPool);
        }

        if (StringUtils.equals(className, "org/springframework/core/type/classreading/SimpleMetadataReader")) {
            return transformSimpleMetadataReader(classPool);
        }

        return null;
    }

    protected byte[] transformSimpleMetadataReader(ClassPool classPool) {
        try {
            CtClass clazz = classPool.get("org.springframework.core.type.classreading.SimpleMetadataReader");
            CtConstructor method = clazz.getDeclaredConstructors()[0];
            String code = getSourceCode(SIMPLE_METADATA_READER_SOURCE_FILE_NAME);
            method.insertBeforeBody(code);
            return clazz.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    protected byte[] transformMainMethodRunner(ClassPool classPool) {
        try {
            CtClass clazz = classPool.get("org.springframework.boot.loader.MainMethodRunner");
            CtMethod method = clazz.getDeclaredMethod("run");
            String code = getSourceCode(MAIN_METHOD_RUNNER_SOURCE_FILE_NAME);
            method.insertBefore(code);
            return clazz.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    private String getSourceCode(String fileName) throws IOException {
        StringBuilder builder = new StringBuilder();

        try (InputStream inputStream = getClass().getResourceAsStream(fileName);
             BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            String line = null;
            while ((line = br.readLine()) != null) {
                builder.append(line).append("\r\n");
            }
        }

        return builder.toString();
    }
}
